"""
    Plant gene-rich regions(GRR) base local density:
        Identifying clusters in plant genomes based
        on non-fixed length gene density.
    input:
        species:
            Species name.
        bed_file:
            chromosome  start   end Transcript_ID   Gene_symbol strand
            JAJAGQ010000025	101605	104317	IMPAAC1M00000012395	IMPAAC1G00000012395	-
            JAJAGQ010000025	120038	125212	IMPAAC1M00000012396	IMPAAC1G00000012396	-
            JAJAGQ010000025	165995	173518	IMPAAC1M00000012397	IMPAAC1G00000012397	+
            JAJAGQ010000025	181051	182061	IMPAAC1M00000012398	IMPAAC1G00000012398	-
            JAJAGQ010000025	243434	245475	IMPAAC1M00000012399	IMPAAC1G00000012399	+
        base_num: 
            The number of bases in the genome.
        gene_num:
            The number of genes in the genome.
        output_table:
            The folder where the file output is located.
    output:
        局部基因密度结果: local density json
        GRR识别结果: 物种___cluster_table_2.xlsx
"""


import os
import json
import copy
import random
import argparse 
import pandas as pd 

from collections import Counter


def parse_bed_to_dict(df_path:str):
    """ 
        input: bed file path
        output: one dict
            chromosome
                |___start
                |   |___region genes information
                |___start sorted list
    """
    bed_df = pd.read_csv(df_path, sep="\t",header=None)
    # BED file format, required input format.
    title = ["chromosome","start","end","Transcript_ID","Gene_symbol","strand"]
    bed_df.columns = title
    # Generate a dictionary mapping Gene_symbol to Transcript_ID.
    gene_id_transcript_id_dict = {}
    for index, row in bed_df.iterrows():
        if row["Gene_symbol"] not in gene_id_transcript_id_dict:
            gene_id_transcript_id_dict[row["Gene_symbol"]] = []
        gene_id_transcript_id_dict[row["Gene_symbol"]].append(row["Transcript_ID"])
    # Deduplicate the table using the values in the Gene_symbol column, 
    # noting that some Gene_symbols may have multiple Transcript_ID rows.
    bed_df = bed_df.drop_duplicates(subset="Gene_symbol", keep="last")
    bed_dict = {}
    # Group by chromosome name.
    chrom_group = bed_df.groupby("chromosome")
    for chrom, chrom_df in chrom_group:
        if chrom not in bed_dict:
            bed_dict[chrom] = {} 
        start_list = []
        for row in chrom_df.iterrows():
            start = row[1]["start"]
            start_list.append(int(start))
            if start not in bed_dict[chrom]:
                bed_dict[chrom][start] = {}
            end = row[1]["end"]
            # This is for writing to and viewing the file.
            key = "____".join([str(start), str(end)]) 
            # key = tuple([start, end])
            value = row[1].tolist() 
            bed_dict[chrom][start][key] = value 
        bed_dict[chrom]["start_list"] = sorted(list(set(start_list)))
    # Output regions with the same start position.
    # for c,c_v in bed_dict.items():
    #     start_list = c_v.pop("start_list")
    #     for s,s_v in c_v.items():
    #         if len(s_v) > 1:
    #             print(s_v)

    return bed_dict


def gene_density_average_3(chrom_value, start_list):
    """
        Iterate through the genes of each chromosome/scaffold/contig, 
        using three genes as a sliding frame.
        Compute the local gene density.
        output: 
            chromosome[
                [sliding_frame_index_list],
                [sliding_frame_density_list]
            ]
    """
    index = 0 
    # Initialize a start_list_copy, which will change as the loop progresses.
    start_list_copy = copy.deepcopy(start_list)  
    sliding_frame_density_list = [""]*(len(start_list_copy) - 2)
    sliding_frame_index_list = [""]*(len(start_list_copy) - 2)
    # print(len(sliding_frame_density_list))
    # print(len(sliding_frame_index_list))
    # print("start_list",len(start_list))
    while len(start_list_copy) > 2 and index < len(sliding_frame_index_list):  
        # The starting position of the gene where the sliding frame begins.
        cluster_start = start_list_copy[0]
        # Initialize end_index since the sliding frame is 
        # configured to require at least three elements.
        # Therefore, the initialization is to obtain the starting position of the third gene.
        end_index = 2
        # Obtain the starting position of the third gene.
        start_3 = start_list_copy[end_index]
        # Obtain the ending position of the third gene.
        end_3_dict = chrom_value[start_3]
        end_3 = ""
        # A starting position may correspond to multiple regions; 
        # the ending position should be chosen as the longest one.
        for start_end, values in end_3_dict.items(): 
            if end_3 == "":
                end_3 = values[2]
            else:
                if values[2] > end_3:
                    end_3 = values[2]
        # Initialize start_end_length, 
        # which is used to calculate the length of the sliding frame.
        cluster_density = (end_3 - cluster_start + 1) / 3 
        # print("index: ", index, cluster_density, len(chrom_value)) 
        cluster_density = round((1/(cluster_density/1000)) * 10, 4)
        if cluster_density < 0:
            print(cluster_density)
        # print(cluster_density, "\n")
        try:
            sliding_frame_density_list[index] = cluster_density
            sliding_frame_index_list[index] = index
            index = index + 1
            start_list_copy = start_list[index:]
            # print(index, len(chrom_value), len(start_list_copy), cluster_density)
        except Exception as e:
            print("\n---------------- erorr ---------------")
            print(index)
            print(e)
            print("---------------- erorr ---------------\n")
        
    return [sliding_frame_index_list, sliding_frame_density_list]


def calculate_gene_3_density(genome_gene_pfam_dict):
    """ 
        Calculate the local gene density for three genes on 
        the genome using a sliding window approach.
    """
    local_density_dict = {}
    for chromosome, chrom_value in genome_gene_pfam_dict.items():
        chrom_value_ = copy.deepcopy(chrom_value)
        start_list = chrom_value_.pop("start_list")
        # Each chromosome/scaffold/contig must have at least 
        # three genes to calculate gene density.
        if len(start_list) >= 3:
            density_result = gene_density_average_3(chrom_value_, start_list)
            if density_result != {}:
                # If density_result is not empty, 
                # add the result to local_density_dict.
                local_density_dict[chromosome] = density_result
        else:
            # print(chromosome, start_list)
            pass
            
    return local_density_dict


def get_density_quartile(local_density_dict):
    """ 
        Calculate the quartiles of local gene density across the genome.
    """
    data_list = []
    for chromosome, chrom_value in local_density_dict.items():
        data_list.extend(chrom_value[1])
    data_series = pd.Series(data_list) 
    # Calculate quartiles.  
    quartiles = data_series.quantile([0.25, 0.5, 0.75]).tolist()   
    statistic_dict = {
        "min_value": float(data_series.min()),
        "Q1": quartiles[0],
        "Q2": quartiles[1],
        "Q3": quartiles[2],
        "max_value": float(data_series.max())
    }
    
    return statistic_dict


def filter_density_dict(local_density_dict, statistic_dict):
    """ 
        Filter out sliding windows with local gene density less than Q3.      
    """
    Q3 = statistic_dict["Q3"]
    density_filter_dict = {}
    for chromosome, chrom_value in local_density_dict.items():
        if chromosome not in density_filter_dict:
            density_filter_dict[chromosome] = []
        frame_index_list = chrom_value[0]
        gene_density_list = chrom_value[1]
        index_dnesity = zip(frame_index_list, gene_density_list)
        for row in index_dnesity:
            if row[1] >= Q3:
                density_filter_dict[chromosome].append(row)
    
    return density_filter_dict
    
    
def extract_continuous_position(filtered_local_density_dict):
    """ 
        Extract sliding windows that are contiguously located after filtering.
    """
    continuous_position_dict = {}
    for chromosome, chrom_value in filtered_local_density_dict.items():
        if chromosome not in continuous_position_dict:
            continuous_position_dict[chromosome] = {}
        index_start = 0
        cluster_name = 0
        while index_start < len(chrom_value):
            tmp_chrom_value = copy.deepcopy(chrom_value[index_start:])
            tmp_cluster = [tmp_chrom_value[0]]
            index_num = 0
            while True:
                start_site = tmp_chrom_value[index_num][0]
                # Exit the loop if the index exceeds 
                # the length of the sliding window list.
                if index_num + 1 == len(tmp_chrom_value):
                    break
                # Condition for discontinuity judgment.
                if tmp_chrom_value[index_num + 1][0] != start_site + 1:
                    break
                tmp_cluster.append(tmp_chrom_value[index_num + 1])
                index_num += 1
                index_start += 1
            continuous_position_dict[chromosome][cluster_name] = tmp_cluster
            cluster_name += 1
            index_start += 1

    return continuous_position_dict
            

def filter_tmp_cluster_dict(continuous_position_dict):
    """ 
        Temporary clusters where the local gene density 
        in three consecutive sliding windows is greater than Q3 are retained.
    """ 
    filter_result = {}
    for chrom,chrom_v in continuous_position_dict.items():
        if chrom not in filter_result:
            filter_result[chrom] = {}
        for k,v in chrom_v.items():
            if len(v) > 0:
                filter_result[chrom][k] = v
    
    return filter_result


def gene_cluster_extend_add_info(gene_cluster, bed_dict):
    """ 
        The gene cluster is extended by two genes and information is added. 
    """
    result = []
    cluster_id = 1
    for chrom, chrom_cluster in gene_cluster.items():
        chrom_info_dict = copy.deepcopy(bed_dict[chrom])
        start_list = chrom_info_dict.pop("start_list")
        for k,sliding_index_densitys in chrom_cluster.items():
            # print(k,sliding_index_densitys)
            start_sliding = sliding_index_densitys[0]
            # print(start_sliding)
            start_sliding_index = start_sliding[0]
            # print(start_sliding_index)
            start_sliding_start = start_list[start_sliding_index]
            # print(start_sliding_start)
            end_sliding = sliding_index_densitys[-1]
            end_sliding_index = end_sliding[0]
            end_sliding_start = start_list[end_sliding_index]
            
            extend_two_rows = []
            end = 0
            # Extend the first gene.
            if start_list[end_sliding_index + 1] in chrom_info_dict:
                infos = chrom_info_dict[start_list[end_sliding_index + 1]]
                for start_e, info in infos.items():
                    row = [
                        chrom, cluster_id, info[3], info[4], info[1], info[2],
                        info[-1]
                    ]
                    extend_two_rows.append(row)
            # Extend the second gene.
            # print(start_list[end_sliding_index + 2], len(start_list))
            if end_sliding_index + 2 < len(start_list):
                infos = chrom_info_dict[start_list[end_sliding_index + 2]]
                for start_e, info in infos.items():
                    row = [
                        chrom, cluster_id, info[3], info[4], info[1], info[2],
                        info[-1]
                    ]
                    if info[2] > end:
                        end = info[2]
                    extend_two_rows.append(row)
    
            cluster_gene_num = (len(sliding_index_densitys)) + len(extend_two_rows)
            cluster_gene_density =  round(
                (1 / (((end - start_sliding_start + 1) / cluster_gene_num) / 1000)) * 10, 4)
            if cluster_gene_density < 0:
                print(chrom, start_sliding_index, cluster_gene_density)
            # print((1 / (((end - end_sliding_start) / 3) / 1000)) * 10)
            # print("\n")
                
            for sliding_index_density in sliding_index_densitys:
                start_index = sliding_index_density[0]
                density = sliding_index_density[1]
                start = start_list[start_index]
                infos = chrom_info_dict[start]
                for start_e, info in infos.items():
                    row = [
                        chrom, cluster_id, info[3], info[4], info[1], info[2],
                        info[-1], density, cluster_gene_density
                    ]
                    result.append(row)
            extend_two_rows = [
                x + ["", cluster_gene_density] for x in extend_two_rows
            ]
            result.extend(extend_two_rows)
            cluster_id += 1
            
    return result
            

def cluster_json_parse_to_table(
    gene_cluster_info, species_name, statistic_dict, genome_info, title):
    """ 
        Generate a results table. 
    """
    data = []
    for row in gene_cluster_info:
        text = (
            [species_name] + genome_info + row + [statistic_dict["min_value"], 
            statistic_dict["Q1"], statistic_dict["Q2"], statistic_dict["Q3"], 
            statistic_dict["max_value"]]
        ) 
        data.append(text)
    df = pd.DataFrame(data, columns=title)
    
    return df 


def merge_cluster_shared_genes(cluster_table):
    """ 
        Merge clusters of shared genes. 
    """
    gene_ids = cluster_table["gene_id"].tolist()
    counter = Counter(gene_ids)
    duplicate_genes = [gene for gene, count in counter.items() if count > 1]
    duplicate_genes_df = cluster_table[
        cluster_table["gene_id"].isin(duplicate_genes)
    ]
    # print("\nThe number of genes appearing in two or more clusters is:\n\t", 
    #       len(duplicate_genes))
    duplicate_cluster = sorted(list(set(duplicate_genes_df["cluster_name"].tolist())))
    # print("\nThe number of gene clusters with overlapping genes is:\n\t",
    #       len(list(set(duplicate_cluster))))
    require_merge_cluster_df = cluster_table[
        cluster_table["cluster_name"].isin(duplicate_cluster)
    ]
    not_require_merge_cluster_df = cluster_table[
        ~cluster_table["cluster_name"].isin(duplicate_cluster)
    ]
    require_merge_cluster_dict = {
        key[0]:sub_df
        for key,sub_df in require_merge_cluster_df.groupby(["cluster_name"])
    }
    # with open("test_require_merge_cluster_dict.json",mode="w") as file:
    #     json.dump(
    #         {k:str(v.shape) for k,v in require_merge_cluster_dict.items()}, 
    #         file, indent=4
    #     )
    merge_result_list = []
    while duplicate_cluster != []:
        tmp_cluster = duplicate_cluster.pop()
        # print(tmp_cluster)
        tmp_cluster_df = require_merge_cluster_dict[tmp_cluster]
        tmp_cluster_genes = tmp_cluster_df["gene_id"].tolist()
        cluster_merge = []
        for cluster in duplicate_cluster:
            if cluster != tmp_cluster:
                cluster_genes = require_merge_cluster_dict[cluster]["gene_id"].tolist()
                for gene in cluster_genes:
                    if gene in tmp_cluster_genes:
                        cluster_merge.append(cluster)
                        # print("\n\n",duplicate_cluster,"\n\n")
                        if cluster in duplicate_cluster:
                            index = duplicate_cluster.index(cluster)
                            duplicate_cluster.pop(index)
        # print(cluster_merge)
        if cluster_merge != []:
            for cluster in cluster_merge:
                merge_df = require_merge_cluster_dict[cluster]
                merge_df["cluster_name"] = [tmp_cluster] * merge_df.shape[0]
                tmp_cluster_df = pd.concat([tmp_cluster_df, merge_df])
                
        tmp_cluster_df = tmp_cluster_df.sort_values(
            by=["start","end"], ascending=[True, True]
        )
        # print(tmp_cluster_df[["start","end"]],"\n")
        start = tmp_cluster_df["start"].tolist()[0]
        end = tmp_cluster_df["end"].tolist()[-1]
        gene_number = len(list(set(tmp_cluster_df["gene_id"].tolist())))
        density = ( 1 / (( (end-start + 1) / gene_number ) / 1000 ) ) * 10
        density = round(density, 4)
        tmp_cluster_df["cluster_gene_density"] = [density] * tmp_cluster_df.shape[0]
        merge_result_list.append(tmp_cluster_df)
    merge_result_df = pd.concat(merge_result_list)
    all_df = pd.concat([not_require_merge_cluster_df, merge_result_df])
    # Reassign cluster numbers
    index = 1
    result_df_list = []
    for cluster, sub_df in all_df.groupby(["cluster_name"]):
        sub_df["cluster_name"] = index
        result_df_list.append(sub_df)
        index += 1
    result_df = pd.concat(result_df_list)
                
    return result_df
    

def plant_gene_cluster_base_local_density(
        species_name, all_base_num, all_gene_num, bed_path, output_table_dir):
    """ Main """
    # Calculate the average gene density across the genome
    length = int(all_base_num)
    gene_number = int(all_gene_num)
    gene_density = length / gene_number
    genome_info = [
        gene_number, length, 
        round((1/(gene_density/1000)) * 10, 4)
    ]
    print(
        "\n01.genome_info:"
        "\n\t The length of the genome is: ", length, "\n\t",
        "The number of genes in the genome is: ", gene_number, "\n\t",
        "The average gene density of the genome is: ", 
        str(round((1/(gene_density/1000)) * 10, 4)) + " genes per 10 kb", "\n"
    )

    # Convert a BED file to a dictionary. bed_dict
    bed_dict = parse_bed_to_dict(bed_path)
    # with open("test_bed_dict.json", mode="a") as file:
    #     json.dump(bed_dict, file, indent=4)
    
    # Calculate the local gene density for three genes 
    # on the genome using a sliding window approach.
    local_density_dict = calculate_gene_3_density(bed_dict)
    os.makedirs("./01.local_density_json/", exist_ok=True)
    with open("./01.local_density_json/"+species_name+".json", mode="w") as file:
        json.dump(local_density_dict, file, indent=4)
    
    # Obtain the quartiles of local gene density across the genome.
    statistic_dict = get_density_quartile(local_density_dict)
    print("\n02.Statistical measures of local gene density across the genome: \n\t", 
        statistic_dict, "\n")
    
    # Filter out sliding windows with local gene density less than Q3.
    filtered_local_density_dict = filter_density_dict(local_density_dict, statistic_dict)
    # with open("test_filtered_local_density_dict.json", mode="w") as file:
    #     json.dump(filtered_local_density_dict, file, indent=4)
    
    # Extract sliding windows that are contiguously located after filtering.
    continuous_position_dict = extract_continuous_position(filtered_local_density_dict)
    # with open("test_continuous_position_dict.json", "w") as file:
    #     json.dump(continuous_position_dict, file, indent=4)
    
    # Temporary clusters where the local gene density 
    # in three consecutive sliding windows is greater than Q3 are retained.
    gene_cluster = filter_tmp_cluster_dict(continuous_position_dict)
    
    title = [
        "species", "genome_gene_number","genome_length","genome_gene_density",
        "chromosome/scaffold/contif","cluster_name", "transcript_id", "gene_id", 
        "start", "end",  "strand", "gene_local_density", "cluster_gene_density", 
        "min","quartile_lower", "quartile_median", "quartile_upper", "max"]
    
    # The gene cluster is extended by two genes and information is added.
    gene_cluster_info = gene_cluster_extend_add_info(gene_cluster, bed_dict)
    
    # Generate a results table.
    cluster_table = cluster_json_parse_to_table(
        gene_cluster_info, species_name, statistic_dict, genome_info, title
    )
    # cluster_table.to_excel(
    #     output_table_path + "/"+ species_name + "___cluster_table_1.xlsx", 
    #     index=False
    # )
    # print(cluster_table.head())
    print("\n03.Number of gene clusters before merging: \n\t", 
        len(list(set(cluster_table["cluster_name"].tolist())))
    )
    
    # Merge clusters of shared genes.
    merge_cluster_result_df = merge_cluster_shared_genes(cluster_table)
    print("\n04.Number of gene clusters after merging: \n\t", 
        len(list(set(merge_cluster_result_df["cluster_name"].tolist()))),
        "\n  The number of genes in the gene cluster is:\n\t",
        merge_cluster_result_df.shape[0], "\n"
    )
    
    merge_cluster_result_df.to_excel(
        output_table_path + "/"+ species_name + "___cluster_table_2.xlsx", 
        index=False
    )


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-species","--species_name")
    parser.add_argument("-base_num","--all_base_num")
    parser.add_argument("-gene_num", "--all_gene_num")
    parser.add_argument("-bed","--bed_file_path")
    parser.add_argument("-output_table","--output_table_path")
    args = parser.parse_args()
    
    return args
    
    
if __name__ == "__main__":
    params = parse_args()
    species_name = params.species_name
    all_base_num = params.all_base_num
    all_gene_num = params.all_gene_num
    bed_path = params.bed_file_path
    output_table_path = params.output_table_path
    print(species_name, all_base_num, all_gene_num, bed_path, output_table_path)
    
    plant_gene_cluster_base_local_density(
        species_name, all_base_num, all_gene_num, bed_path, output_table_path
    )